#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2020 wanghuacheng <wanghuacheng@wanghuacheng-PC>
#
# Distributed under terms of the MIT license.

""" """
from six import StringIO
import prettytable
import os
import sys
import hks
import pandas as pd

output = StringIO()

import tushare as ts
from datetime import datetime
today = datetime.today().strftime("%Y-%m-%d")
def get_s50():
    sh50 = pd.read_csv('sh50.csv')
    sh50_w = {}
    for index, row in sh50.iterrows():
        code = row['code'].split('.')[0]
        sh50_w[code] = row['weight']
    return sh50_w

def get_a50():
    a50 = pd.read_csv('a50.csv', dtype = {'code': str})
    a50_w = {}
    for index, row in a50.iterrows():
        code = row['code'].split('.')[0]
        a50_w[code] = row['weight']
    return a50_w



def notify(title, content):
    cmd = "notify-send '%s' '%s'" % (str(title), str(content))
    print(cmd)
    os.system(cmd)

def strip(s):
   return s.strip()

code_file = "code.txt"
if len(sys.argv) == 2:
    code_file = sys.argv[1]


with open(code_file) as cf:
   codes = list(map(strip, cf.readlines()))



def format_for_print(df):
    table = prettytable.PrettyTable(list(df.columns))
    for row in df.itertuples():
        table.add_row(row[1:])
    return str(table)


def alert(df, code, limit):
    code_df = df[df.code == code]
    c = code_df.change.values[0]
    name = code_df.name.values[0]
    if c > limit:
        notify(name+":"+code,'change:' + str(c))

def sweight(row, sh50_w):
    w = sh50_w[row.code]
    rw = (row.price / row.pre_close - 1) * w * 100
    return round(rw , 2)

def show_w(index_func):
    sh50_w = index_func()
    fields = ['code', 'name', 'cp',  'pre_close', 'price', 'bid']
    n_fields = ['open', 'price', 'bid', 'pre_close', 'volume'] #, 'v_ma5',  'v_ma10', 'v_ma20']
    wcodes = list(sh50_w.keys())
    df = ts.get_realtime_quotes(wcodes)
    df[n_fields] = df[n_fields].astype('float')
    df['change'] = (df.price / df.pre_close - 1) * 100
    #df['w'] = df.apply(lambda x: (x['price'] / x['pre_close'] -1 ) * sh50_w[x['code']])
    df['w'] = df.apply(lambda row:sweight(row, sh50_w), axis = 1)
    df['h'] = df.code.apply(lambda x:sh50_w.get(x, -1))
    #print(df)
    #df = df[df.w.abs() > 1].sort_values(['w'], ascending = False)
    df = df.sort_values(['w'], ascending = False)
    df = df.round({'change': 1})
    print(format_for_print(df[['code', 'name', 'price', 'change', 'w', 'h']]))



def show_sts():
    #fields = ['code', 'name', 'open',  'pre_close', 'price', 'bid',  'volume', 'change', 'change_percent']
    fields = ['code', 'name', 'cp',  'pre_close', 'price', 'bid']
    n_fields = ['open', 'price', 'bid', 'pre_close', 'volume'] #, 'v_ma5',  'v_ma10', 'v_ma20']
    df = ts.get_realtime_quotes(codes)
    df[n_fields] = df[n_fields].astype('float')
    df['change'] = (df.price / df.pre_close - 1)
    df['cp'] = df['change'].apply(lambda x:format(x, '.2%'))
    print(format_for_print(df[fields]))


def show_index():
    index = ts.get_index()
    #print(index)
    fields = ['code', 'name', 'change', 'close', 'high', 'low', 'preclose']
    print(format_for_print(index[index.code.isin(['000001', '000009', '000010','000016', '000300', '399006', '000905' ])][fields]))

def show_hk():
    table = prettytable.PrettyTable(['name', 'code', 'last_close', 'today', 'rate'])
    codes = ["00700-HK", "01810-HK", "800000-HK", "800700-HK", "03690-HK", "09626-HK", "09988-HK"]
    for c in codes:
        row = hks.get_realtime_st(c)
        if row:
            table.add_row(row)
    print(table)


def main():
    itype = 'st'
    if len(sys.argv) == 3:
        itype = sys.argv[2]
    if itype == 'index':
        show_index()
        show_w(get_s50)
    if itype == 'aindex':
        show_w(get_a50)
    else:
        show_sts()
        show_hk()

if __name__ == '__main__':
    main()
